import torch.nn as nn
from abc import ABCMeta, abstractmethod

class BaseDenseHead(nn.Module, metaclass=ABCMeta):
    """Base class for DenseHeads."""

    def __init__(self):
        super(BaseDenseHead, self).__init__()

    @abstractmethod
    def loss(self, **kwargs):
        """Compute losses of the head."""
        pass

    @abstractmethod
    def predict(self, **kwargs):
        """Compute losses of the head."""
        pass

    def forward_train(self, x, listed_gt_bboxes, listed_labels):
        model_output = self(x)
        losses = self.loss(model_output,listed_gt_bboxes, listed_labels, gt_bboxes_ignore=None)
        return losses

    def forward_val(self, x, listed_gt_bboxes, listed_labels):
        model_output = self(x)
        losses = self.loss(model_output,listed_gt_bboxes, listed_labels, gt_bboxes_ignore=None)
        pred_boxes = self.predict(model_output,listed_gt_bboxes, listed_labels, gt_bboxes_ignore=None)
        return losses, pred_boxes

